{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "sys.path.append(os.path.abspath('../'))\n",
    "from common import *\n",
    "import configs\n",
    "\n",
    "import json\n",
    "\n",
    "fig_name = 'forwardgrads'\n",
    "fig_dir = join(FIGURE_DIR, fig_name)\n",
    "used_configs = [configs.ConvolutionWarp2(), configs.ConvolutionWarp4(),\n",
    "                configs.ConvolutionWarp8(), configs.ConvolutionWarp16(),\n",
    "                configs.ConvolutionWarp32(), configs.Warp(), configs.FiniteDifferences()]\n",
    "\n",
    "scenes = ['plane_area', 'plane_red_object', 'glossy_plane']\n",
    "sdfs = ['shadowing_128', 'logo_256', 'bunny_128']\n",
    "names_pretty = ['Shadowing', 'Logo', 'Bunny']\n",
    "params = ['x', 'x', 'x']\n",
    "exposures = [0, 0, 0.5]\n",
    "w = 128\n",
    "h = 128\n",
    "crop_t = 15\n",
    "crop_b = 15\n",
    "\n",
    "arrows = [\n",
    "    (0.95, 0.5, -0.1, 0.0),\n",
    "    (0.9, 0.18, -0.07, 0.07),\n",
    "    (0.49, 0.47, -0.07, 0.02)\n",
    "]\n",
    "\n",
    "# Timings in forward mode aren't representative, so we don't plot them by default\n",
    "plot_times = False\n",
    "aspect = w / h\n",
    "fontsize = 14\n",
    "base_size = 2\n",
    "n_rows = len(scenes)\n",
    "n_cols = len(params) * len(used_configs) + 1\n",
    "fig = plt.figure(1, figsize=(n_cols * base_size * aspect, n_rows * base_size * 0.78), constrained_layout=False)\n",
    "gs = fig.add_gridspec(n_rows, n_cols, wspace=0.025, hspace=0.025)\n",
    "names = ['Convolution \\n 2 samples', 'Convolution \\n 4 samples', 'Convolution \\n 8 samples',\n",
    "         'Convolution \\n 16 samples', 'Convolution \\n 32 samples', r'\\textbf{Ours}', 'Reference (FD)']\n",
    "y_offset = -0.14\n",
    "for row, (scene, sdf, param, exposure) in enumerate(zip(scenes, sdfs, params, exposures)):\n",
    "    sdf = sdf.split('_')[0]\n",
    "    ax = fig.add_subplot(gs[row, 0])\n",
    "    img = read_img(join(fig_dir, f'{sdf}.exr'), exposure=exposure)\n",
    "    img = img[crop_t:-crop_b, :, :]\n",
    "    img = np.repeat(np.repeat(img, 4, axis=1), 4, axis=0)\n",
    "    aspect = img.shape[0] / img.shape[1]\n",
    "    ax.imshow(img, interpolation='none', extent=[0, 1, 0, aspect])\n",
    "    disable_ticks(ax)\n",
    "    ax.set_ylabel('\\\\textsc{' + names_pretty[row] + '}', labelpad=5)\n",
    "    ax.yaxis.set_label_position(\"left\")\n",
    "    if row == n_rows - 1:\n",
    "        txt = ax.set_title('Scene', fontsize=fontsize, y=y_offset, va='top')\n",
    "    col = 1\n",
    "    r = None\n",
    "    for cfg_idx, cfg in enumerate(used_configs):\n",
    "        if type(cfg) is configs.FiniteDifferences:\n",
    "            cfg.pretty_name_short = 'Reference'\n",
    "\n",
    "        ax = fig.add_subplot(gs[row, col])\n",
    "        fn = join(fig_dir, f'{sdf}_{cfg.name}_{param}.exr')\n",
    "        img = read_img(fn, 5, False)\n",
    "        img = img[crop_t:-crop_b, :, :]\n",
    "        img = np.mean(img, axis=-1)\n",
    "        if r is None:\n",
    "            r = np.quantile(np.abs(img), 0.99)\n",
    "            r = np.maximum(r, 1)\n",
    "\n",
    "        img = np.repeat(np.repeat(img, 4, axis=1), 4, axis=0)\n",
    "        aspect = img.shape[0] / img.shape[1]\n",
    "        ax.imshow(img, vmin=-r, vmax=r, cmap='coolwarm_r', extent=[0, 1, 0, aspect], interpolation='none')\n",
    "        disable_ticks(ax)\n",
    "        if row == n_rows - 1 and col > 0:\n",
    "            txt = ax.set_title(names[col - 1], fontsize=fontsize, y=y_offset, va='top')\n",
    "\n",
    "        if not type(cfg) is configs.FiniteDifferences and plot_times:\n",
    "            with open(join(fig_dir, f'{sdf}_{cfg.name}_{param}.json'), 'r') as f:\n",
    "                stats = json.load(f)\n",
    "            duration = stats['total_time'] / 1000\n",
    "            txt = ax.text(0.99, 0.0, f\"{duration:.2f} s\", size=14, ha=\"right\", va=\"bottom\")\n",
    "            txt.set_path_effects([path_effects.withStroke(linewidth=1.0, foreground='white')])\n",
    "        col += 1\n",
    "        if cfg_idx == 1:\n",
    "            ax.arrow(arrows[row][0], arrows[row][1], arrows[row][2], arrows[row][3], lw=0.75,\n",
    "                     head_width=0.03, edgecolor='k', facecolor='k', overhang=0.1)\n",
    "plt.margins(0, 0)\n",
    "# save_fig(fig_name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('mi')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c8cc192903caa17681aa39d71202092ac11a526d37a1c4ad2948f13605924304"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
